import asyncio
import json
from utils.asynccmd import AsyncCmdMixin
from export.export import export_to_excel
from config.settings import PACKAGE_PATH
import random
import os


def filter_url(x): return x['url']


class DNSXScan(AsyncCmdMixin):
    def __init__(self, path, domains: list, query=''):
        self._exec_path = path
        self.domains = domains
        self.query = query
        # 存放存活的域名
        self.results = []

    def parse(self, results):
        outputs = []
        query_domains = self.domains
        for result in results:
            if 'a' not in result:
                continue
            d = {}
            for a in result['a']:
                d['子域名'] = result['host']
                d['解析结果'] = a
                d['解析类型'] = 'A'
                outputs.append(d)
            query_domains.remove(result['host'])

        if query_domains:
            for unknown in query_domains:
                d = {}
                d['子域名'] = unknown
                d['解析结果'] = ''
                d['解析类型'] = ''
                outputs.append(d)
        return outputs

    async def scan(self, arguments=""):
        """TODO: 修改为批量执行"""
        results = []
        file_name = ''.join(random.sample('zyxwvutsrqponmlkjihgfedcba', 15))
        file_path = PACKAGE_PATH / f'tmp/{file_name}'
        with open(file_path, 'w+') as f:
            f.write('\n'.join(self.domains))
        cmd = f'{self._exec_path} -l {file_path} -silent -json'
        result = await self.run_cmd(cmd)
        if not result:
            self.results = []
            os.remove(file_path)
            return
        for x in result:
            try:
                results.append(json.loads(x))
            except Exception as e:
                continue
        self.results = self.parse(results)
        os.remove(file_path)

    def export(self, is_save, save_path=None):
        if not is_save:
            return
        export_to_excel(self.results, sheet_name=f'domain-{self.query}', xlsx_name='test',
                        save_path=save_path)
